{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Importing fastai will also import numpy, pytorch, etc. \n",
    "from fastai.vision import *\n",
    "from numbers import Integral\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Modified from \n",
    "K=2\n",
    "class MultiTransformLabelList(LabelList):\n",
    "    def __getitem__(self,idxs:Union[int,np.ndarray])->'LabelList':\n",
    "        \"return a single (x, y) if `idxs` is an integer or a new `LabelList` object if `idxs` is a range.\"\n",
    "        idxs = try_int(idxs)\n",
    "        if isinstance(idxs, Integral):\n",
    "            if self.item is None: x,y = self.x[idxs],self.y[idxs]\n",
    "            else:                 x,y = self.item   ,0\n",
    "            if self.tfms or self.tfmargs:\n",
    "                #I've changed this line to return a list of augmented images\n",
    "                x = [x.apply_tfms(self.tfms, **self.tfmargs) for _ in range(K)]\n",
    "            if hasattr(self, 'tfms_y') and self.tfm_y and self.item is None:\n",
    "                y = y.apply_tfms(self.tfms_y, **{**self.tfmargs_y, 'do_resolve':False})\n",
    "            if y is None: y=0\n",
    "            return x,y\n",
    "        else: return self.new(self.x[idxs], self.y[idxs])\n",
    "        \n",
    "#I'll also need to change the default collate function to accomodate multiple augments\n",
    "def MixmatchCollate(batch):\n",
    "    batch = to_data(batch)\n",
    "    if isinstance(batch[0][0],list):\n",
    "        batch = [[torch.stack(s[0]),s[1]] for s in batch]\n",
    "    return torch.utils.data.dataloader.default_collate(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Grab file path to cifar dataset. Will download data if not present\n",
    "path = untar_data(URLs.CIFAR)\n",
    "\n",
    "#Custom ImageList with filter function\n",
    "class MixMatchImageList(ImageList):\n",
    "    def filter_train(self,num_items,seed=2343):\n",
    "        train_idxs = np.array([i for i,o in enumerate(self.items) if Path(o).parts[-3] != \"test\"])\n",
    "        valid_idxs = np.array([i for i,o in enumerate(self.items) if Path(o).parts[-3] == \"test\"])\n",
    "        np.random.seed(seed)\n",
    "        keep_idxs = np.random.choice(train_idxs,num_items,replace=False)\n",
    "        self.items = np.array([o for i,o in enumerate(self.items) if i in np.concatenate([keep_idxs,valid_idxs])])\n",
    "        return self\n",
    "    \n",
    "#Create two databunch objects for the labeled and unlabled images. A fastai databunch is a container for train, validation, and\n",
    "#test dataloaders which automatically processes transforms and puts the data on the gpu.\n",
    "data_labeled = (MixMatchImageList.from_folder(path)\n",
    "                .filter_train(500) #Use 500 labeled images for traning\n",
    "                .split_by_folder(valid=\"test\") #test on all 10000 images in test set\n",
    "                .label_from_folder()\n",
    "                .transform(get_transforms(),size=32)\n",
    "                #On windows, must set num_workers=0. Otherwise, remove the argument for a potential performance improvement\n",
    "                .databunch(bs=64,num_workers=0)\n",
    "                .normalize(cifar_stats))\n",
    "\n",
    "train_set = set(data_labeled.train_ds.x.items)\n",
    "src = (ImageList.from_folder(path)\n",
    "        .filter_by_func(lambda x: x not in train_set)\n",
    "        .split_by_folder(valid=\"test\"))\n",
    "src.train._label_list = MultiTransformLabelList\n",
    "data_unlabeled = (src.label_from_folder()\n",
    "         .transform(get_transforms(),size=32)\n",
    "         .databunch(bs=128,collate_fn=MixmatchCollate,num_workers=0)\n",
    "         .normalize(cifar_stats))\n",
    "\n",
    "#Databunch with all 50k images labeled, for baseline\n",
    "data_full = (ImageList.from_folder(path)\n",
    "        .split_by_folder(valid=\"test\")\n",
    "        .label_from_folder()\n",
    "        .transform(get_transforms(),size=32)\n",
    "        .databunch(bs=128,num_workers=0)\n",
    "        .normalize(cifar_stats))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mixup(a_x,a_y,b_x,b_y,alpha=0.75):\n",
    "    l = np.random.beta(alpha,alpha)\n",
    "    l = max(l,1-l)\n",
    "    x = l * a_x + (1-l) * b_x\n",
    "    y = l* a_y + (1-l) * b_y\n",
    "    return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sharpen(p,T=0.5):\n",
    "    u = p ** (1/T)\n",
    "    return u / u.sum(dim=1,keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MixupLoss(nn.Module):\n",
    "    def forward(self, preds, target, unsort=None, ramp=None, bs=None):\n",
    "        if unsort is None:\n",
    "            return F.cross_entropy(preds,target)\n",
    "        preds = preds[unsort]\n",
    "        preds_l = preds[:bs]\n",
    "        preds_ul = preds[bs:]\n",
    "        preds_l = torch.log_softmax(preds_l,dim=1)\n",
    "        preds_ul = torch.softmax(preds_ul,dim=1)\n",
    "        loss_x = -(preds_l * target[:bs]).sum(dim=1).mean()\n",
    "        loss_u = F.mse_loss(preds_ul,target[bs:])\n",
    "        self.loss_x = loss_x.item()\n",
    "        self.loss_u = loss_u.item()\n",
    "        return loss_x + 100 * ramp * loss_u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MixMatchTrainer(LearnerCallback):\n",
    "    _order=-20\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.l_dl = iter(data_labeled.train_dl)\n",
    "        self.smoothL, self.smoothUL = SmoothenValue(0.98), SmoothenValue(0.98)\n",
    "        self.recorder.add_metric_names([\"l_loss\",\"ul_loss\"])\n",
    "        self.it = 0\n",
    "        \n",
    "    def on_batch_begin(self, train, last_input, last_target, **kwargs):\n",
    "        if not train: return\n",
    "        try:\n",
    "            x_l,y_l = next(self.l_dl)\n",
    "        except:\n",
    "            self.l_dl = iter(data_labeled.train_dl)\n",
    "            x_l,y_l = next(self.l_dl)\n",
    "            \n",
    "        x_ul = last_input\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            ul_labels = sharpen(torch.softmax(torch.stack([self.learn.model(x_ul[:,i]) for i in range(x_ul.shape[1])],dim=1),dim=2).mean(dim=1))\n",
    "            \n",
    "        x_ul = torch.cat([x for x in x_ul])\n",
    "        ul_labels = torch.cat([y.unsqueeze(0).expand(K,-1) for y in ul_labels])\n",
    "        \n",
    "        l_labels = torch.eye(data_labeled.c).cuda()[y_l]\n",
    "        \n",
    "        w_x = torch.cat([x_l,x_ul])\n",
    "        w_y = torch.cat([l_labels,ul_labels])\n",
    "        idxs = torch.randperm(w_x.shape[0])\n",
    "        \n",
    "        mixed_input, mixed_target = mixup(w_x,w_y,w_x[idxs],w_y[idxs])\n",
    "        bn_idxs = torch.randperm(mixed_input.shape[0])\n",
    "        unsort = [0] * len(bn_idxs)\n",
    "        for i,j in enumerate(bn_idxs): unsort[j] = i\n",
    "        mixed_input = mixed_input[bn_idxs]\n",
    "    \n",
    "\n",
    "        ramp = self.it / 3000.0 if self.it < 3000 else 1.0\n",
    "        return {\"last_input\": mixed_input, \"last_target\": (mixed_target,unsort,ramp,x_l.shape[0])}\n",
    "    \n",
    "    def on_batch_end(self, train, **kwargs):\n",
    "        if not train: return\n",
    "        self.smoothL.add_value(self.learn.loss_func.loss_x)\n",
    "        self.smoothUL.add_value(self.learn.loss_func.loss_u)\n",
    "        self.it += 1\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics,[self.smoothL.smooth,self.smoothUL.smooth])    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EMATrainer(LearnerCallback):\n",
    "    _order=-20\n",
    "    def __init__(self, learn, ema_model):\n",
    "        super().__init__(learn)\n",
    "        self.ema_model = ema_model.cuda()\n",
    "        self.ema_model.load_state_dict(self.learn.model.state_dict())\n",
    "        \n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.recorder.add_metric_names([\"ema_acc\"])\n",
    "        \n",
    "    def on_batch_end(self, train, **kwargs):\n",
    "        if not train: return\n",
    "        a = 0.999\n",
    "        model_sd = self.learn.model.state_dict()\n",
    "        ema_sd = self.ema_model.state_dict()\n",
    "        for key in ema_sd.keys():\n",
    "            if key.find(\".bn.\") >= 0:\n",
    "                ema_sd[key].data.copy_(model_sd[key].data)\n",
    "            else:\n",
    "                ema_sd[key].data.mul_(a).add_((1-a)*model_sd[key].data).mul_(1-(2e-3*2e-2))\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        preds, labels = [],[]\n",
    "        self.ema_model.eval()\n",
    "        with torch.no_grad():\n",
    "            for x,y in iter(self.learn.data.valid_dl):\n",
    "                p = self.ema_model(x)\n",
    "                preds.append(torch.argmax(p,dim=1))\n",
    "                labels.append(y)\n",
    "        preds = torch.cat(preds)\n",
    "        labels = torch.cat(labels)\n",
    "        acc = (preds == labels).float().mean().item() \n",
    "        return add_metrics(last_metrics,[acc])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.WideResNet(num_groups=3,N=4,num_classes=10,k=2,start_nf=32)\n",
    "ema_model = models.WideResNet(num_groups=3,N=4,num_classes=10,k=2,start_nf=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = Learner(data_unlabeled,model,loss_func=MixupLoss(),callback_fns=[MixMatchTrainer],metrics=accuracy)\n",
    "learn.callbacks.append(EMATrainer(learn,ema_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(200,2e-3,wd=0.02)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
